import argparse
import os
from models.classifier import UltrasoundClassifier
from models.ultrasam_wrapper import UltraSamWrapper

def main(image_path):
    if not os.path.exists(image_path):
        print("Error: image not found:", image_path)
        return

    # 1) 分类
    clf = UltrasoundClassifier()
    organ = clf.predict(image_path)
    print(f"🧠 识别结果：{organ}")

    # 2) UltraSam 分割（占位/包装）
    sam = UltraSamWrapper()
    mask_path = image_path + "_mask.png"
    sam.infer_to_file(image_path, mask_path)
    print("✅ 分割结果已保存：", mask_path)
    print("说明：如果你已放入真实 UltraSam 权重并替换 wrapper 中的加载逻辑，分割将为真实结果。")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", required=True, help="输入超声图像路径（PNG/JPG/…）")
    args = parser.parse_args()
    main(args.image)
